/*
 *@Project:
 *@Author: wanghong
 *@Date: 2019年07月09日
 *@Copyright: ©2019-2028 www.xfxb.net Inc. All rights reserved.
 */
package com.dopper.ws;

import com.dopper.ws.annotation.WebsocketMapping;
import com.dopper.ws.common.ThreadFactoryImpl;
import com.dopper.ws.exception.PathMappingException;
import io.netty.bootstrap.ServerBootstrap;
import io.netty.buffer.PooledByteBufAllocator;
import io.netty.channel.*;
import io.netty.channel.epoll.EpollEventLoopGroup;
import io.netty.channel.epoll.EpollServerSocketChannel;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.handler.codec.http.HttpObjectAggregator;
import io.netty.handler.codec.http.HttpServerCodec;
import io.netty.handler.ssl.*;
import io.netty.handler.stream.ChunkedWriteHandler;
import io.netty.handler.timeout.IdleStateHandler;
import io.netty.util.concurrent.DefaultEventExecutorGroup;
import io.netty.util.concurrent.EventExecutorGroup;
import io.netty.util.internal.StringUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.net.ssl.SSLException;
import java.io.File;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author wanghong
 * @date 2019年07月09日
 */
public class WebsocketServer {

    private final static Logger log = LoggerFactory.getLogger(WebsocketServer.class);

    private final static String PATH_STARTS_WITH = "/";

    private Set<Websocket> websockets = new HashSet<Websocket>();

    private Map<String, Websocket> websocketMap = new HashMap<String,Websocket>();

    private Map<Websocket, WebsocketMsgHandler> websocketMsgHandlerMap = new HashMap<Websocket, WebsocketMsgHandler>();

    private EventLoopGroup bossGroup;

    private EventLoopGroup workerGroup;

    private WebsocketServerConfig config;

    private EventExecutorGroup workThreadPool;

    public WebsocketServer() {
        this(new WebsocketServerConfig());
    }

    public WebsocketServer(WebsocketServerConfig config) {
        this.config = config;
    }

    public void registerWebsocket(Websocket websocket) {
        websockets.add(websocket);
    }

    Websocket mapping(String path) {
        return websocketMap.get(path);
    }

    WebsocketMsgHandler mappingHandler(Websocket websocket) {
        return websocketMsgHandlerMap.get(websocket);
    }

    EventExecutorGroup getWorkThreadPool() {
        return workThreadPool;
    }

    public void run(String host,int port) throws Exception {
        if (StringUtil.isNullOrEmpty(host)) {
            host = "0.0.0.0";
        }
        final WebsocketServer server = this;

        initNettyGroup();

        final SslContext finalSslContext = loadSslContext();

        try {

            ServerBootstrap b = new ServerBootstrap();
            b.group(bossGroup, workerGroup)
                .channel(WebsocketServerConfig.isIsLinuxPlatform() ? EpollServerSocketChannel.class : NioServerSocketChannel.class)
                .option(ChannelOption.SO_BACKLOG, config.getBacklog())
                .option(ChannelOption.SO_REUSEADDR, true)
                //.option(ChannelOption.SO_KEEPALIVE, false)
                .childOption(ChannelOption.TCP_NODELAY, true)
                .childOption(ChannelOption.SO_SNDBUF, config.getServerSocketSndBufSize())
                .childOption(ChannelOption.SO_RCVBUF, config.getServerSocketRcvBufSize())
                .childHandler(new ChannelInitializer<SocketChannel>() {

                    @Override
                    protected void initChannel(SocketChannel ch)
                        throws Exception {
                        ChannelPipeline pipeline = ch.pipeline();
                        if (finalSslContext != null) {
                            pipeline.addFirst("sslHandler", finalSslContext.newHandler(ch.alloc()));
                        }
                        pipeline.addLast("http-codec",
                            new HttpServerCodec());
                        pipeline.addLast("aggregator",
                            new HttpObjectAggregator(65536));
                        ch.pipeline().addLast("http-chunked",
                            new ChunkedWriteHandler());
                        if (config.getServerChannelMaxIdleTimeSeconds() > 0) {
                            pipeline.addLast(new IdleStateHandler(0, 0, config.getServerChannelMaxIdleTimeSeconds()));
                        }
                        pipeline.addLast( workThreadPool ,"ws-handler",
                            new WebSocketServerHandler(server));
                    }
                });

            if (this.config.isServerPooledByteBufAllocatorEnable()) {
                b.childOption(ChannelOption.ALLOCATOR, PooledByteBufAllocator.DEFAULT);
            }

            Channel ch = b.bind(host, port).sync().channel();

            handleMapping();

            log.info("{} server started success on {}:{}", config.isUseSsl() ? "wss" : "ws", host, port);
            Runtime.getRuntime().addShutdownHook(new Thread(){
                @Override
                public void run() {
                    close();
                }
            });

            ch.closeFuture().sync();
        } finally {
            bossGroup.shutdownGracefully();
            workerGroup.shutdownGracefully();
            workThreadPool.shutdownGracefully();
        }
    }

    private void initNettyGroup() {
        workThreadPool = new DefaultEventExecutorGroup(config.getServerWorkerThreads(), new ThreadFactoryImpl("WorkThread_"));
        if (config.useEpoll()) {
            this.bossGroup = new EpollEventLoopGroup(config.getServerBossThreads(), new ThreadFactory() {
                private AtomicInteger threadIndex = new AtomicInteger(0);

                @Override
                public Thread newThread(Runnable r) {
                    return new Thread(r, String.format("NettyEPOLLBoss_%d", this.threadIndex.incrementAndGet()));
                }
            });

            this.workerGroup = new EpollEventLoopGroup(config.getServerSelectorThreads(), new ThreadFactory() {
                private AtomicInteger threadIndex = new AtomicInteger(0);
                private int threadTotal = config.getServerSelectorThreads();

                @Override
                public Thread newThread(Runnable r) {
                    return new Thread(r, String.format("NettyServerEPOLLSelector_%d_%d", threadTotal, this.threadIndex.incrementAndGet()));
                }
            });
        } else {
            this.bossGroup = new NioEventLoopGroup(config.getServerBossThreads(), new ThreadFactory() {
                private AtomicInteger threadIndex = new AtomicInteger(0);

                @Override
                public Thread newThread(Runnable r) {
                    return new Thread(r, String.format("NettyNIOBoss_%d", this.threadIndex.incrementAndGet()));
                }
            });

            this.workerGroup = new NioEventLoopGroup(config.getServerSelectorThreads(), new ThreadFactory() {
                private AtomicInteger threadIndex = new AtomicInteger(0);
                private int threadTotal = config.getServerSelectorThreads();

                @Override
                public Thread newThread(Runnable r) {
                    return new Thread(r, String.format("NettyServerNIOSelector_%d_%d", threadTotal, this.threadIndex.incrementAndGet()));
                }
            });
        }
    }

    private SslContext loadSslContext() {
        SslContext sslContext = null;
        if (config.isUseSsl()) {
            SslProvider provider;
            if (OpenSsl.isAvailable()) {
                provider = SslProvider.OPENSSL;
                log.info("Using OpenSSL provider");
            } else {
                provider = SslProvider.JDK;
                log.info("Using JDK SSL provider");
            }
            SslContextBuilder sslContextBuilder = SslContextBuilder
                    .forServer(new File(config.getSslCertPath()), new File(config.getSslKeyPath()), config.getSslKeyPassword())
                    .sslProvider(provider).clientAuth(ClientAuth.NONE);
            try {
                sslContext = sslContextBuilder.build();
            } catch (SSLException e) {
                log.error("初始化ssl时异常", e);
                close();
            }
        }
        return sslContext;
    }

    private void handleMapping() {
        for (Websocket websocket : websockets) {
            WebsocketMapping annotation = websocket.getClass().getAnnotation(WebsocketMapping.class);
            String mapping = PATH_STARTS_WITH;
            if (annotation != null) {
                mapping = annotation.value();
            }
            if (!mapping.startsWith(PATH_STARTS_WITH)) {
                mapping = PATH_STARTS_WITH + mapping;
            }
            if (websocketMap.containsKey(mapping)) {
                log.error("ws服务器启动失败,请求的路径已经存在已经存在 path= [{} ]", mapping);
                throw new PathMappingException(mapping);
            }
            websocketMap.put(mapping, websocket);
            log.info("ws mapping: [{}]", mapping);
            websocketMsgHandlerMap.put(websocket, new WebsocketMsgHandler(websocket));
        }
        if (websocketMap.size() == 0) {
            log.error("no mapping and server close");
            close();
        }
    }

    public void close() {
        bossGroup.shutdownGracefully();
        workerGroup.shutdownGracefully();
        workThreadPool.shutdownGracefully();
    }
}
